#!/usr/bin/python3
# SPDX-License-Identifier: LGPL-3.0-or-later
# Copyright (C) 2023 Intel Corporation
#                    Borys Popławski <borysp@invisiblethingslab.com>
#                    Wojtek Porczyk <woju@invisiblethingslab.com>

import datetime
import re
import sys
import textwrap
import typing

import click

from graminelibos import (
    Manifest, get_tbssigstruct, SGX_LIBPAL,
)

# TODO: after python (>= 3.10) simplify this
# NOTE: we can't `try: importlib.metadata`, because the API has changed between 3.9 and 3.10
# (in 3.9 and in backported importlib_metadata entry_points() doesn't accept group argument)
if sys.version_info >= (3, 10):
    from importlib.metadata import entry_points # pylint: disable=import-error,no-name-in-module
else:
    from pkg_resources import iter_entry_points as entry_points

def list_sgx_sign_plugins():
    return tuple(ep.name for ep in entry_points(group='gramine.sgx_sign'))

_sgx_sign_plugins = list_sgx_sign_plugins()

def get_sgx_sign_plugin(name):
    for ep in entry_points(group='gramine.sgx_sign'):
        if ep.name == name:
            return ep.load()
    raise KeyError(name)

class BCDDate(typing.NamedTuple):
    """Date-compatible object that supports invalid dates (like 0000-00-00)
    """
    year: int
    month: int
    day: int

class BCDDateParamType(click.ParamType):
    """Parse "today" or a date-like string YYYY-MM-DD, accepting invalid dates.

    This is needed to support dates like 0000-00-00, which are invalid days, but valid things to put
    into SIGSTRUCT. Arbitrary values that have hex digits A-F are not currently supported.
    """
    name = 'date'
    _regexp = re.compile(r'(?P<year>\d{4})-(?P<month>\d{2})-(?P<day>\d{2})')
    def convert(self, value, param, ctx):
        if value.casefold() == 'today':
            value = datetime.date.today().strftime('%Y-%m-%d')
        match = self._regexp.fullmatch(value)
        if not match:
            self.fail('expecting "today" or YYYY-MM-DD')
        return BCDDate(**{k: int(v) for k, v in match.groupdict().items()})

@click.command(
    context_settings={'ignore_unknown_options': True},
    epilog=textwrap.dedent(f'''
        Use --with=PLUGIN --help-PLUGIN to get help about particular plugin.

        Available plugins: {", ".join(_sgx_sign_plugins)}'''),
)
@click.option('--with', 'with_', metavar='PLUGIN',
    type=click.Choice(_sgx_sign_plugins),
    default='file',
    help='Choose plugin with which to sign the enclave (default: file)')
@click.option('--output', '-o',
    type=click.Path(),
    help='Output .manifest.sgx file (manifest augmented with autogenerated fields)')
@click.option('--libpal', '-l',
    type=click.Path(exists=True, dir_okay=False),
    default=SGX_LIBPAL,
    help='Input libpal file')
@click.option('--manifest', '-m', 'manifest_file',
    type=click.File('r', encoding='utf-8'),
    help='Input .manifest file')
@click.option('--date', '-d',
    type=BCDDateParamType(), default='today',
    help='Set DATE field in SIGSTRUCT to this value (YYYY-MM-DD or "today")')
@click.option('--chroot',
    type=click.Path(exists=True, dir_okay=True, file_okay=False),
    help='Measure a chroot directory, not the host filesystem')
@click.option('--sigfile', '-s',
    help='Output .sig file')
@click.option('--depfile',
    type=click.File('w'),
    help='Generate dependencies for .manifest.sgx and .sig files')
@click.option('--verbose/--quiet', '-v/-q',
    default=True,
    help='Display details (on by default)')
@click.argument('plugin_args',
    nargs=-1,
    type=click.UNPROCESSED)
@click.pass_context
def main(ctx, with_, output, libpal, manifest_file, date, sigfile, depfile, verbose, plugin_args,
         chroot):
    # pylint: disable=too-many-arguments, too-many-locals

    ret = get_sgx_sign_plugin(with_)(args=plugin_args, standalone_mode=False)

    if isinstance(ret, int):
        # What happened:
        # - user wrote --with=PLUGIN --help-PLUGIN,
        # - subcommand's help_option called ctx.exit() eagerly
        # - Context.exit() raised click.Exit exception
        # - this exception was caught in Context.main() and handled depending on standalone_mode
        #   and in standalone_mode=False, it just returned e.exit_code (which is probably 0).
        # Therefore, we also exit with the same exit_code.
        ctx.exit(ret)

    if output is None:
        ctx.fail('Missing option --output')
    if manifest_file is None:
        ctx.fail('Missing option --manifest')

    try:
        it = iter(ret)
        # no TypeError, therefore we've got tuple or list, or sth else iterable
        sign_func, extra_deps = it
    except TypeError:
        # sign_func is probably just a callable (no need to check, will break later if it isn't)
        # and extra dependencies were not provided
        sign_func, extra_deps = ret, ()

    manifest = Manifest.load(manifest_file)

    try:
        expanded = manifest.expand_all_trusted_files(chroot=chroot)
    except FileNotFoundError as err:
        ctx.fail(f'Missing trusted file: {err.filename!r}')

    with open(output, 'wb') as f:
        manifest.dump(f)

    if not sigfile:
        if manifest_file.name.endswith('.manifest'):
            sigfile = manifest_file.name[:-len('.manifest')]
        else:
            sigfile = manifest_file.name
        sigfile += '.sig'

    sigstruct = get_tbssigstruct(output, date, libpal, verbose=verbose)
    sigstruct.sign(sign_func)

    with open(sigfile, 'wb') as f:
        f.write(sigstruct.to_bytes())

    if depfile:
        # Dependencies:
        #
        # - `.manifest.sgx` depends on all files we just expanded
        # - `.sig` additionally depends on libpal
        #
        # TODO (Ninja 1.10): We print all these as dependencies for `.manifest.sgx`. This will still
        # cause `.sig` to be rebuilt when necessary: we build both these files together, so it's not
        # possible to rebuild one without the other.
        #
        # This is a workaround for the fact that Ninja prior to version 1.10 does not
        # support depfiles with multiple outputs (and parses such depfiles incorrectly).
        deps = [*expanded, libpal, *extra_deps]

        depfile.write(f'{output}:')
        for filename in deps:
            depfile.write(f' \\\n\t{filename}')
        depfile.write('\n')


if __name__ == '__main__':
    main() # pylint: disable=no-value-for-parameter
